package main

import (
	"encoding/json"
	"errors"
	"os"
	"path/filepath"
	"sync"
)

type UpstreamServer struct {
	// target like "http://127.0.0.1:9000" or "https://10.0.0.2:9443"
	Target string `json:"target"`
	// weight for weighted round-robin (1-100, default: 1)
	Weight int `json:"weight,omitempty"`
}

type Upstream struct {
	// Single target (for backward compatibility)
	Target string `json:"target,omitempty"`
	// Multiple servers for load balancing
	Servers []UpstreamServer `json:"servers,omitempty"`
	// Load balancing algorithm: "round-robin" or "weight" (default: "round-robin")
	Algorithm string `json:"algorithm,omitempty"`
}

type DomainConfig struct {
	Domain   string   `json:"domain"`
	Upstream Upstream `json:"upstream"`
	// optional certificate files for TLS termination on :443
	CertFile string `json:"cert_file,omitempty"`
	KeyFile  string `json:"key_file,omitempty"`
}

type FileConfig struct {
	// List of domain configs
	Domains []DomainConfig `json:"domains"`
}

// State holds in-memory routing and TLS certificate handles
type State struct {
	mu sync.RWMutex

	// domain -> upstream
	routes map[string]Upstream

	// domain -> cert file paths (for reload on demand)
	certPaths map[string][2]string
}

func NewState() *State {
	return &State{
		routes:    make(map[string]Upstream, 32),
		certPaths: make(map[string][2]string, 32),
	}
}

func LoadStateFromFile(path string) (*State, error) {
	st := NewState()
	data, err := os.ReadFile(path)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			// create empty file config
			if err := persistFileConfig(path, &FileConfig{Domains: []DomainConfig{}}); err != nil {
				return nil, err
			}
			return st, nil
		}
		return nil, err
	}
	var fc FileConfig
	if len(data) > 0 {
		if err := json.Unmarshal(data, &fc); err != nil {
			return nil, err
		}
	}
	st.mu.Lock()
	defer st.mu.Unlock()
	for _, dc := range fc.Domains {
		up := dc.Upstream
		// 兼容旧格式：如果只有 Target，转换为 Servers 格式
		if up.Target != "" && len(up.Servers) == 0 {
			up.Servers = []UpstreamServer{{Target: up.Target, Weight: 1}}
			up.Algorithm = "round-robin"
		}
		// 设置默认算法
		if up.Algorithm == "" {
			up.Algorithm = "round-robin"
		}
		// 设置默认权重
		for i := range up.Servers {
			if up.Servers[i].Weight <= 0 {
				up.Servers[i].Weight = 1
			}
		}
		st.routes[dc.Domain] = up
		if dc.CertFile != "" && dc.KeyFile != "" {
			st.certPaths[dc.Domain] = [2]string{dc.CertFile, dc.KeyFile}
		}
	}
	return st, nil
}

func (s *State) Snapshot() FileConfig {
	s.mu.RLock()
	defer s.mu.RUnlock()
	out := FileConfig{Domains: make([]DomainConfig, 0, len(s.routes))}
	for d, up := range s.routes {
		dc := DomainConfig{Domain: d, Upstream: up}
		// 如果只有一个服务器且使用轮询，转换为旧格式以便兼容
		if len(up.Servers) == 1 && up.Algorithm == "round-robin" && up.Target == "" {
			dc.Upstream.Target = up.Servers[0].Target
			dc.Upstream.Servers = nil
		}
		if ck, ok := s.certPaths[d]; ok {
			dc.CertFile = ck[0]
			dc.KeyFile = ck[1]
		}
		out.Domains = append(out.Domains, dc)
	}
	return out
}

func (s *State) SetRoute(domain string, up Upstream) {
	s.mu.Lock()
	// 如果只有 Target，转换为 Servers 格式
	if up.Target != "" && len(up.Servers) == 0 {
		up.Servers = []UpstreamServer{{Target: up.Target, Weight: 1}}
		up.Algorithm = "round-robin"
	}
	// 设置默认算法
	if up.Algorithm == "" {
		up.Algorithm = "round-robin"
	}
	// 设置默认权重
	for i := range up.Servers {
		if up.Servers[i].Weight <= 0 {
			up.Servers[i].Weight = 1
		}
	}
	s.routes[domain] = up
	s.mu.Unlock()
}

func (s *State) DeleteRoute(domain string) {
	s.mu.Lock()
	delete(s.routes, domain)
	delete(s.certPaths, domain)
	s.mu.Unlock()
}

func (s *State) GetUpstream(host string) (Upstream, bool) {
	s.mu.RLock()
	up, ok := s.routes[stripPort(host)]
	s.mu.RUnlock()
	return up, ok
}

func (s *State) SetCertPath(domain, certFile, keyFile string) {
	s.mu.Lock()
	s.certPaths[domain] = [2]string{certFile, keyFile}
	s.mu.Unlock()
}

func (s *State) GetCertPaths() map[string][2]string {
	s.mu.RLock()
	defer s.mu.RUnlock()
	out := make(map[string][2]string, len(s.certPaths))
	for k, v := range s.certPaths {
		out[k] = v
	}
	return out
}

func persistFileConfig(path string, fc *FileConfig) error {
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return err
	}
	data, err := json.MarshalIndent(fc, "", "  ")
	if err != nil {
		return err
	}
	tmp := path + ".tmp"
	if err := os.WriteFile(tmp, data, 0o644); err != nil {
		return err
	}
	return os.Rename(tmp, path)
}

func stripPort(hostport string) string {
	for i := len(hostport) - 1; i >= 0; i-- {
		if hostport[i] == ':' {
			// IPv6 bracket form [::1]:443 — keep brackets by finding rightmost ':'
			return hostport[:i]
		}
	}
	return hostport
}
